use netneighbours;
use pcap::{Capture, Device};
use std::net::IpAddr;

use std::ptr;

use libc::c_char;
use std::ffi::CStr;

type MacAddr = [u8; 6];
type Ipv4Addr = [u8; 4];

#[repr(C)]
pub struct Addresses {
    pub source_mac_addr: MacAddr,
    pub source_ip_addr: Ipv4Addr,
    pub des_mac_addr: MacAddr,
    pub des_ip_addr: Ipv4Addr,
}

#[repr(C)]
pub enum None_or<T> {
    Some(T),
    None,
}

#[no_mangle]
pub extern "C" fn get_default_addrs() -> None_or<Addresses> {
    match get_addrs() {
        Ok(addrs) => None_or::Some(addrs),
        Err(e) => {
            println!("Failed to get network interface: {}", e);
            None_or::None
        }
    }
}

// 获得一个可用的网卡的名称作为默认网卡名称。
#[no_mangle]
pub extern "C" fn find_default_dev_name() -> None_or<*mut c_char> {
    // 首先调用 get_addrs() 获得可用网卡的IP地址
    // get_addrs() 使用了 ipconfig 库，该库可以获得本机的IP地址、MAC地址、网关地址、DNS服务器地址等信息。
    // 通过检查 ipconfig 库找到的网卡的网关地址来确定此网卡是否可用。
    let addresses = match get_addrs() {
        Ok(addresses) => addresses,
        Err(e) => {
            eprintln!("Failed to get network interface: {}", e);
            return None_or::None;
        }
    };

    // 调用 pcap 库的 Device::list() 函数获得所有的网卡设备。
    let devices = match Device::list() {
        Ok(devices) => devices,
        Err(e) => {
            eprintln!("Failed to list devices: {}", e);
            return None_or::None;
        }
    };

    // 然后遍历所有的网卡设备，找到与本机IP地址匹配的网卡设备。
    // 如果找到了，则返回该网卡的名称。
    for dev in devices {
        if dev.addresses.iter().any(|addrs| {
            if let IpAddr::V4(addr) = addrs.addr {
                addr.octets() == addresses.source_ip_addr
            } else {
                false
            }
        }) {
            let name = dev.name;
            let str_len = name.len();
            let src = name.as_bytes().as_ptr();
            let c_name = unsafe {
                let malloc = libc::malloc(str_len + 1) as *mut u8;
                if malloc.is_null() {
                    ptr::null_mut()
                } else {
                    ptr::copy_nonoverlapping(src, malloc, str_len);
                    let cursor = malloc.add(str_len);
                    ptr::write(cursor, 0);
                    malloc as *mut c_char
                }
            };
            return None_or::Some(c_name);
        }
    }
    None_or::None
}


#[no_mangle]
pub extern "C" fn send_dns_packet(
    pkt_id: u16,
    s_port: u16,
    addrs: *const Addresses,
    domain: *const c_char,
) -> u8 {
    if addrs.is_null() || domain.is_null() {
        eprintln!("Invalid input: null pointer found.");
        return 1; // 返回错误代码
    }

    // 安全地解引用addrs指针
    let addrs = unsafe { &*addrs };

    let domain_name = unsafe {
        match CStr::from_ptr(domain).to_str() {
            Ok(s) => s,
            Err(_) => {
                eprintln!("Failed to convert C string.");
                return 1; // 返回错误代码
            }
        }
    };

    if let Err(e) = send_dns(pkt_id, s_port, addrs, domain_name) {
        println!("Send packet faild: {}", e);
        return 1;
    }
    return 0;
}

// 获得本机MAC地址、路由器MAC地址、本机IP地址，DNS服务器IP地址。
pub fn get_addrs() -> std::result::Result<Addresses, String> {
    match ipconfig::get_adapters() {
        Ok(adapters) => {
            let available_adapter = adapters
                .iter()
                .find(|adapter| !adapter.gateways().is_empty())
                .ok_or("No available adapter found")?;

            let mut mac_addr = [0u8; 6];
            mac_addr.copy_from_slice(
                &available_adapter
                    .physical_address()
                    .ok_or("No mac address found")?,
            );

            let ip_addr = available_adapter
                .ip_addresses()
                .iter()
                .find_map(|addr| {
                    if let IpAddr::V4(v4) = addr {
                        Some(v4.octets())
                    } else {
                        None
                    }
                })
                .ok_or("No ip address found.")?;

            let router_ip_addr = available_adapter
                .gateways()
                .iter()
                .find(|&addr| addr.is_ipv4())
                .ok_or("No router mac address found")?;

            let mut router_mac_addr = [0u8; 6];
            router_mac_addr.copy_from_slice(
                &netneighbours::get_table()
                    .iter()
                    .find(|x| x.0 == *router_ip_addr)
                    .ok_or("No router mac address found")?
                    .1
                    .as_bytes(),
            );

            let dns_server = available_adapter
                .dns_servers()
                .iter()
                .find_map(|addr| {
                    if let IpAddr::V4(v4) = addr {
                        Some(v4.octets())
                    } else {
                        None
                    }
                })
                .ok_or("No dns server found")?;

            return Ok(Addresses {
                source_mac_addr: mac_addr,
                source_ip_addr: ip_addr,
                des_mac_addr: router_mac_addr,
                des_ip_addr: dns_server,
            });
        }
        Err(_) => return Err("No adapter found".to_string()),
    };
}

pub fn send_dns(
    pkt_id: u16,
    s_port: u16,
    info: &Addresses,
    domain: &str,
) -> std::result::Result<(), String> {
    let device = Device::list()
        .map_err(|e| format!("Failed to list devices: {}", e))?
        .into_iter()
        .find(|dev| {
            dev.addresses.iter().any(|addrs| {
                if let IpAddr::V4(v4addr) = addrs.addr {
                    v4addr.octets() == info.source_ip_addr
                } else {
                    false
                }
            })
        })
        .ok_or("No device found")?;

    let mut cap = Capture::from_device(device).unwrap().open().unwrap();

    // Ethernet frame header
    let mut ethernet_frame = [
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // Destination MAC address
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // Source MAC address
        0x08, 0x00, // Type field (0x0800 for IPv4)
    ];

    ethernet_frame[0..6].copy_from_slice(&info.des_mac_addr);
    ethernet_frame[6..12].copy_from_slice(&info.source_mac_addr);

    let mut ip_packet = [
        0x45, // Version: 4, Header Length: 5 (20 bytes)
        0x00, // Differentiated Services Field: 0
        0x00, 0x00, // Total Length: Will be filled later
        0x00, 0x00, // Identification: 0
        0x40, 0x00, // Flags: Don't Fragment, Fragment Offset: 0
        0x80, // Time to Live: 128
        0x11, // Protocol: UDP (17)
        0x00, 0x00, // Header Checksum: Will be filled later
        0x00, 0x00, 0x00, 0x00, // Source IP Address: Will be filled later
        0x00, 0x00, 0x00, 0x00, // Destination IP Address: Will be filled later
    ];

    ip_packet[12..16].copy_from_slice(&info.source_ip_addr);

    ip_packet[16..20].copy_from_slice(&info.des_ip_addr);

    // UDP Header
    let mut udp_packet = [
        0x00, 0x00, // Source Port: Will be filled later
        0x00, 0x35, // Destination Port: 53
        0x00, 0x00, // Length: Will be filled later
        0x00, 0x00, // Checksum: Will be filled later
    ];

    // fill source Port
    udp_packet[0] = (s_port >> 8) as u8;
    udp_packet[1] = (s_port & 0xFF) as u8;

    // DNS Header
    let mut dns_packet = [
        0x00, 0x00, // Identifier: Will be filled later
        0x01, 0x00, // Flags: Standard Query
        0x00, 0x01, // Questions: 1
        0x00, 0x00, // Answer RRs: 0
        0x00, 0x00, // Authority RRs: 0
        0x00, 0x00, // Additional RRs: 0
    ];

    dns_packet[0] = (pkt_id >> 8) as u8;
    dns_packet[1] = (pkt_id & 0xFF) as u8;

    // DNS Question
    let mut question = Vec::new();
    for part in domain.split('.') {
        question.push(part.len() as u8);
        question.extend_from_slice(part.as_bytes());
    }

    question.push(0); // End of domain name

    let question_type = [0x00, 0x01]; // Type A (Host Address)
    let question_class = [0x00, 0x01]; // Class IN (Internet)

    // Calculate and fill the udp packet length
    let udp_length = udp_packet.len()
        + dns_packet.len()
        + question.len()
        + question_type.len()
        + question_class.len();

    udp_packet[4] = (udp_length >> 8) as u8;
    udp_packet[5] = (udp_length & 0xFF) as u8;

    // Calculate and fill the udp packet checksum
    let pseudo_header = [
        ip_packet[12],
        ip_packet[13],
        ip_packet[14],
        ip_packet[15], // Source IP Address
        ip_packet[16],
        ip_packet[17],
        ip_packet[18],
        ip_packet[19], // Destination IP Address
        0x00,          // Zero
        ip_packet[9],  // Protocol
        udp_packet[4],
        udp_packet[5], // UDP Length
    ];

    // Concatenate dns header, question, question type and question class
    let mut dns_payload = Vec::new();
    dns_payload.extend_from_slice(&dns_packet);
    dns_payload.extend_from_slice(&question);
    dns_payload.extend_from_slice(&question_type);
    dns_payload.extend_from_slice(&question_class);

    let checksum = checksum(&pseudo_header, &udp_packet, &dns_payload);

    udp_packet[6] = (checksum >> 8) as u8;
    udp_packet[7] = (checksum & 0xFF) as u8;

    // Calculate and fill the IP packet total length
    let ip_length = ip_packet.len() + udp_length;
    ip_packet[2] = (ip_length >> 8) as u8;
    ip_packet[3] = (ip_length & 0xFF) as u8;

    // Calculate and fill the IP packet header checksum
    let checksum = ip_checksum(&ip_packet);
    ip_packet[10] = (checksum >> 8) as u8;
    ip_packet[11] = (checksum & 0xFF) as u8;

    // Create a buffer for the full packet
    let mut packet = Vec::new();

    // Add the Ethernet frame to the buffer
    packet.extend_from_slice(&ethernet_frame);

    // Add the IP packet to the buffer
    packet.extend_from_slice(&ip_packet);

    // Add the UDP packet to the buffer
    packet.extend_from_slice(&udp_packet);

    // Add the DNS packet to the buffer
    packet.extend_from_slice(&dns_payload);

    // Send the packet
    match cap.sendpacket(packet) {
        Ok(_) => {
            println!("Packet sent successfully.");
            Ok(())
        }
        Err(e) => Err(format!("Failed to send packet: {}", e)),
    }
}
pub fn send_icmp(info: Addresses) -> std::result::Result<(), String> {
    let device = Device::list()
        .map_err(|e| format!("Failed to list devices: {}", e))?
        .into_iter()
        .find(|dev| {
            dev.addresses.iter().any(|addrs| {
                if let IpAddr::V4(v4addr) = addrs.addr {
                    v4addr.octets() == info.source_ip_addr
                } else {
                    false
                }
            })
        })
        .ok_or("No device found")?;

    let mut cap = Capture::from_device(device).unwrap().open().unwrap();

    // Ethernet frame header
    let mut ethernet_frame = [
        0x10, 0x4f, 0x58, 0x6c, 0x0c, 0x00, // Destination MAC address
        0x60, 0xe9, 0xaa, 0x5c, 0x00, 0x5b, // Source MAC address
        0x08, 0x00, // Type field (0x0800 for IPv4)
    ];

    ethernet_frame[0..6].copy_from_slice(&info.des_mac_addr);
    ethernet_frame[6..12].copy_from_slice(&info.source_mac_addr);

    let mut ip_packet = [
        0x45, // Version: 4, Header Length: 5 (20 bytes)
        0x00, // Differentiated Services Field: 0
        0x00, 0x00, // Total Length: Will be filled later
        0x00, 0x00, // Identification: 0
        0x40, 0x00, // Flags: Don't Fragment, Fragment Offset: 0
        0x80, // Time to Live: 128
        0x01, // Protocol: ICMP (1)
        0x00, 0x00, // Header Checksum: Will be filled later
        0x00, 0x00, 0x00, 0x00, // Source IP Address: Will be filled later
        0x00, 0x00, 0x00, 0x00, // Destination IP Address: Will be filled later
    ];

    ip_packet[12..16].copy_from_slice(&info.source_ip_addr);

    ip_packet[16..20].copy_from_slice(&info.des_ip_addr);

    // let destination_ip = Ipv4Addr::new(60, 204, 241, 149); // Replace with your destination IP
    // ip_packet[16..20].copy_from_slice(&destination_ip.octets());

    // ICMP Echo Request Header
    let mut icmp_packet = [
        0x08, // Type: 8 (Echo Request)
        0x00, // Code: 0
        0x00, 0x00, // Checksum: Will be filled later
        0x00, 0x00, // Identifier: 0
        0x00, 0x01, // Sequence Number: 1
        // Payload
        0x48, 0x65, 0x6c, 0x6c, 0x6f, 0x2c, 0x20, 0x77, 0x6f, 0x72, 0x6c, 0x64,
        0x21, // "Hello, world!"
    ];

    // Calculate and fill the icmp checksum
    let checksum = icmp_checksum(&icmp_packet);
    icmp_packet[2] = (checksum >> 8) as u8;
    icmp_packet[3] = (checksum & 0xFF) as u8;

    // Calculate and fill the ip packet total length
    let total_length = ip_packet.len() + icmp_packet.len();
    ip_packet[2] = (total_length >> 8) as u8;
    ip_packet[3] = (total_length & 0xFF) as u8;

    // Calculate and fill the ip packet header checksum
    let checksum = ip_checksum(&ip_packet);
    ip_packet[10] = (checksum >> 8) as u8;
    ip_packet[11] = (checksum & 0xFF) as u8;

    // Create a buffer for the full packet
    let mut packet = Vec::new();

    // Add the Ethernet frame to the buffer
    packet.extend_from_slice(&ethernet_frame);

    // Add the IP packet to the buffer
    packet.extend_from_slice(&ip_packet);

    // Add the ICMP packet to the buffer
    packet.extend_from_slice(&icmp_packet);

    // Send the packet
    match cap.sendpacket(packet) {
        Ok(_) => {
            println!("Packet sent successfully.");
            Ok(())
        }
        Err(e) => Err(format!("Failed to send packet: {}", e)),
    }
}

pub fn send_tcp_syn(
    s_port: u16,
    info: &Addresses,
    seq: u32,
    ack: u32
) -> std::result::Result<Capture<pcap::Active>, String> {
    let device = Device::list()
        .map_err(|e| format!("Failed to list devices: {}", e))?
        .into_iter()
        .find(|dev| {
            dev.addresses.iter().any(|addrs| {
                if let IpAddr::V4(v4addr) = addrs.addr {
                    v4addr.octets() == info.source_ip_addr
                } else {
                    false
                }
            })
        })
        .ok_or("No device found")?;

    let mut cap: Capture<pcap::Active> = Capture::from_device(device).unwrap().open().unwrap();

    // Ethernet frame header
    let mut ethernet_frame = [
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // Destination MAC address
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // Source MAC address
        0x08, 0x00, // Type field (0x0800 for IPv4)
    ];

    ethernet_frame[0..6].copy_from_slice(&info.des_mac_addr);
    ethernet_frame[6..12].copy_from_slice(&info.source_mac_addr);

    // IP packet
    let mut ip_packet = [
        0x45, // Version: 4, Header Length: 5 (20 bytes)
        0x00, // Differentiated Services Field: 0
        0x00, 0x00, // Total Length: Will be filled later
        0x00, 0x00, // Identification: 0
        0x40, 0x00, // Flags: Don't Fragment, Fragment Offset: 0
        0x80, // Time to Live: 128
        0x06, // Protocol: TCP (6)
        0x00, 0x00, // Header Checksum: Will be filled later
        0x00, 0x00, 0x00, 0x00, // Source IP Address: Will be filled later
        0x00, 0x00, 0x00, 0x00, // Destination IP Address: Will be filled later
    ];

    ip_packet[12..16].copy_from_slice(&info.source_ip_addr);
    ip_packet[16..20].copy_from_slice(&info.des_ip_addr);

    // TCP packet
    let mut tcp_packet = [
        0x00, 0x00, // Source Port: Will be filled later
        // 0x00, 0x50, // Destination Port: 80
        0x1F, 0x42, // Destination Port: 8002
        0x00, 0x00, 0x00, 0x00, // Sequence Number: 0
        0x00, 0x00, 0x00, 0x00, // Acknowledgement Number: 0
        0x50, // Data Offset: 5 (20 bytes), Reserved: 0
        0x02, // Flags: SYN
        0xFF, 0xFF, // Window Size: 65535
        0x00, 0x00, // Checksum: Will be filled later
        0x00, 0x00, // Urgent Pointer: 0
    ];

    // Fill the sequence number and acknowledgement number field of the TCP packet
    tcp_packet[4..8].copy_from_slice(&seq.to_be_bytes());
    tcp_packet[8..12].copy_from_slice(&ack.to_be_bytes());

    // Fill the source port field of the TCP packet
    tcp_packet[0..2].copy_from_slice(&s_port.to_be_bytes());

    // Calculate the checksum for the TCP packet
    // TCP pseudo header
    let mut pseudo_header = [
        0x00, 0x00, 0x00, 0x00, // Source IP Address: Will be filled later
        0x00, 0x00, 0x00, 0x00, // Destination IP Address: Will be filled later
        0x00, // Zero
        0x06, // Protocol: TCP (6)
        0x00, 0x14, // TCP Length: Will be filled later
    ];
    
    pseudo_header[0..4].copy_from_slice(&info.source_ip_addr);
    pseudo_header[4..8].copy_from_slice(&info.des_ip_addr);
    let tcp_length = tcp_packet.len() as u16;
    pseudo_header[10..12].copy_from_slice(&tcp_length.to_be_bytes());
    
    // Calculate the checksum for the TCP packet
    let checksum = checksum(&pseudo_header, &tcp_packet, &[]);
    tcp_packet[16..18].copy_from_slice(&checksum.to_be_bytes());

    // Fill the total length field of the IP packet
    let total_length = (ip_packet.len() + tcp_packet.len()) as u16;
    ip_packet[2..4].copy_from_slice(&total_length.to_be_bytes());
    
    // Calculate the checksum for the IP packet
    let checksum = ip_checksum(&ip_packet);
    ip_packet[10..12].copy_from_slice(&checksum.to_be_bytes());

    // Send the TCP packet
    let mut packet = Vec::new();
    packet.extend_from_slice(&ethernet_frame);
    packet.extend_from_slice(&ip_packet);
    packet.extend_from_slice(&tcp_packet);

    match cap.sendpacket(packet) {
        Ok(_) => {
            println!("Packet sent successfully.");
            Ok(cap)
        }
        Err(e) => Err(format!("Failed to send packet: {}", e)),
    }
}

fn icmp_checksum(packet: &[u8]) -> u16 {
    let mut sum = 0u32;
    let mut i = 0;

    while i + 1 < packet.len() {
        let word = (packet[i] as u32) << 8 | packet[i + 1] as u32;
        sum = sum + word;
        i = i + 2;
    }

    // If the packet length is odd, add the last byte with padding
    if i < packet.len() {
        let word = (packet[i] as u32) << 8;
        sum = sum + word;
    }

    while (sum >> 16) > 0 {
        sum = (sum & 0xFFFF) + (sum >> 16);
    }

    return !(sum as u16);
}

fn ip_checksum(packet: &[u8]) -> u16 {
    let mut sum = 0u32;
    let mut i = 0;

    while i < packet.len() {
        let word = (packet[i] as u32) << 8 | packet[i + 1] as u32;
        sum = sum + word;
        i = i + 2;
    }

    while (sum >> 16) > 0 {
        sum = (sum & 0xFFFF) + (sum >> 16);
    }

    return !sum as u16;
}

fn checksum(pseudo_header: &[u8], header: &[u8], payload: &[u8]) -> u16 {
    let mut sum = 0u32;

    // Add pseudo header
    for words in pseudo_header.chunks(2) {
        let mut word = [0u8; 2];
        word.copy_from_slice(words);
        sum += u16::from_be_bytes(word) as u32;
    }

    // Add header
    for words in header.chunks(2) {
        let mut word = [0u8; 2];
        word.copy_from_slice(words);
        sum += u16::from_be_bytes(word) as u32;
    }

    // Add payload
    let mut payload = payload.to_vec();
    if payload.len() % 2 != 0 {
        payload.push(0); // Padding if payload length is odd
    }
    for words in payload.chunks(2) {
        let mut word = [0u8; 2];
        word.copy_from_slice(words);
        sum += u16::from_be_bytes(word) as u32;
    }

    // Add overflow
    while (sum >> 16) > 0 {
        sum = (sum & 0xFFFF) + (sum >> 16);
    }

    // Return complement
    !sum as u16
}
